import numpy as np
from scipy.linalg import svd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def fit_line_3d(points):
    """
    该函数用于根据一系列 3D 点拟合直线
    :param points: 输入的 3D 点，形状为 (n, 3) 的 numpy 数组
    :return: 直线上的一个点 p0 和直线的方向向量 v
    """
    # 计算点的均值
    p0 = np.mean(points, axis=0)
    # 将点中心化
    centered_points = points - p0
    # 计算协方差矩阵
    cov_matrix = np.cov(centered_points.T)
    # 使用奇异值分解 (SVD) 求解协方差矩阵的特征向量和特征值
    U, S, Vt = svd(cov_matrix)
    # 最小特征值对应的特征向量就是直线的方向向量
    v = Vt[2]
    return p0, v


def plot_3d_line_and_points(points, p0, v):
    """
    该函数用于绘制 3D 点和拟合的直线
    :param points: 输入的 3D 点，形状为 (n, 3) 的 numpy 数组
    :param p0: 直线上的一个点
    :param v: 直线的方向向量
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # 绘制原始点
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], color='b', label='Original Points')
    # 生成直线上的点
    t = np.linspace(-10, 10, 100)
    line_points = p0 + t[:, np.newaxis] * v
    # 绘制直线
    ax.plot(line_points[:, 0], line_points[:, 1], line_points[:, 2], color='r', label='Fitted Line')
    ax.legend()
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.show()


# 示例数据
points = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7]])


# 拟合直线
p0, v = fit_line_3d(points)


# 绘制 3D 点和拟合的直线
plot_3d_line_and_points(points, p0, v)
